import os, shutil
import numpy as np


def fold_split():
    portion = 0.2
    txt_dir = r"/Users/edy/Downloads/hand_label/yolo_txt"
    image_dir = r"/Users/edy/Downloads/hand_label/Images"

    txt_file = list(filter(lambda x: x.endswith('.txt'), os.listdir(txt_dir)))
    np.random.shuffle(txt_file)
    valid_1 = txt_file[:int(len(txt_file) * portion)]
    train_1 = txt_file[int(len(txt_file) * portion):]

    for txt in valid_1:
        # txt 存到指定文件夹
        shutil.copy((os.path.join(txt_dir, txt)),
                    os.path.join("/Users/edy/Downloads/hand_label/fold_1/train/labels/val", txt))
        img_name = txt.replace('.txt', '.jpg')
        shutil.copy((os.path.join(image_dir, img_name)),
                    os.path.join("/Users/edy/Downloads/hand_label/fold_1/train/images/val", img_name))

    for txt in train_1:
        # txt 存到指定文件夹
        shutil.copy((os.path.join(txt_dir, txt)),
                    os.path.join("/Users/edy/Downloads/hand_label/fold_1/train/labels/train", txt))
        img_name = txt.replace('.txt', '.jpg')
        shutil.copy((os.path.join(image_dir, img_name)),
                    os.path.join("/Users/edy/Downloads/hand_label/fold_1/train/images/train", img_name))


def check_split():
    val_img_path = r"/Users/edy/Downloads/hand_label/fold_1/train/images/val"
    train_img_path = r"/Users/edy/Downloads/hand_label/fold_1/train/images/train"
    train_imgs = os.listdir(train_img_path)
    for img in os.listdir(val_img_path):
        if img in train_imgs:
            raise ("split failed")


fold_split()
# check_split()
